from fastapi import FastAPI, Request, Response, Form
from fastapi.responses import HTMLResponse, RedirectResponse
from cryptography.fernet import Fernet
import bcrypt

users = {
    101: {
        "name": "张三",
        "title": "软件工程师",
    },
    102: {
        "name": "李四",
        "title": "主任",
    },
}

user_passwords = {
    "zhangsan": {
        "user_sn": 101,
        "passwd": "$2b$12$62J.t12HKaanNigybmY7uu5hoQtkiWu5OdpHkltFEx4/NvNdRHBGa",
    },
    "lisi": {
        "user_sn": 102,
        "passwd": "$2b$12$fLh7eaC/WzgNk6/WPv6o0uZo2Q.RkSAKy0ps2t5K9HNG4fr81GiV2",
    },
}

app = FastAPI()

# 生成加密密钥
secret_key = b"GLK064KitB9hnf5rtEWvyO61-BdspfO6Us-Q5WVR4_A="
fernet = Fernet(secret_key)


@app.get("/", response_class=HTMLResponse)
async def main(request: Request, response: Response):
    user_sn = get_current_user(request)
    if not user_sn:
        return RedirectResponse(url="/login", status_code=302)

    user_sn = int(user_sn)
    user = users.get(user_sn)

    html_content = f"""
    <html><body>
    <h3> 您好, {user["title"]} {user["name"]}</h3>

    <form action="/logout" method="post">
        <input type="submit" value="退出登录">
    </form>
    </body></html>
    """
    return HTMLResponse(content=html_content, status_code=200)


@app.get("/login", response_class=HTMLResponse)
async def login_form_page():
    html_content = """
    <html><body>
    <form action="/login" method="post">
        UserName: <input type="text" name="username">
        Password: <input type="password" name="password">
        <input type="submit" value="Sign in">
    </form>
    </body></html>
    """
    return HTMLResponse(content=html_content)


@app.post("/login")
async def login_action(username: str = Form(...), password: str = Form(...)):
    if session_id := check_passwd(username, password):
        response = RedirectResponse(url="/", status_code=302)
        set_current_user(response, session_id)
        return response

    return RedirectResponse(url="/login", status_code=302)


@app.post("/logout")
async def handle_logout(response: Response):
    set_current_user(response, None)
    return RedirectResponse(url="/login", status_code=302)


def check_passwd(username, passwd) -> str:
    principal = user_passwords.get(username)
    if principal is None:
        return

    hashed = principal.get("passwd")
    if not hashed:
        return

    passwd, hashed = passwd.encode("utf-8"), hashed.encode("utf-8")
    if bcrypt.checkpw(passwd, hashed):
        session_id = principal.get("user_sn")
        return f"{session_id}"

    return


def get_current_user(request: Request):
    return get_secure_cookie(request, "session_id")


def set_current_user(response: Response, session_id: str):
    """设置当前会话，若session_id位`None`则清除当前会话."""

    if session_id is not None:
        set_secure_cookie(response, "session_id", session_id)
    else:
        response.delete_cookie("session_id")


def get_secure_cookie(request: Request, name):
    value = request.cookies.get(name)
    if value is None:
        return None

    try:
        buffer = value.encode()
        secured_value = fernet.decrypt(buffer).decode()
        return secured_value
    except Exception:
        print("Cannot decrypt cookie value")
        return None


def set_secure_cookie(response: Response, name, value, **kwargs):
    value = fernet.encrypt(value.encode()).decode()
    response.set_cookie(name, value, **kwargs)
